import os
import time
import re
import numpy as np
import json
from tqdm import tqdm
from glob import glob
import cv2
from PIL import ImageFont, Image, ImageDraw

def draw_pic(image_c, text, img_name01):
    for xx in text:
        i = np.array(xx[1]).reshape(-1).astype(int)
        cv2.line(image_c, (i[0], i[1]), (i[2], i[3]), (0, 0, 255), 2)
        cv2.line(image_c, (i[0], i[1]), (i[6], i[7]), (0, 0, 255), 2)
        cv2.line(image_c, (i[4], i[5]), (i[2], i[3]), (0, 0, 255), 2)
        cv2.line(image_c, (i[4], i[5]), (i[6], i[7]), (0, 0, 255), 2)
    hh, ww, cc = image_c.shape
    max_leng = max(hh, ww, cc)
    if max_leng > 2000:
        offset = min(hh // 100, ww // 100)
    else:
        offset = (hh + ww) // 100
    font = ImageFont.truetype('./files/simhei.ttf', offset, encoding="utf-8")
    img_PIL = Image.fromarray(cv2.cvtColor(image_c, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_PIL)
    for xx in text:
        s = xx[0]
        i = np.array(xx[1]).reshape(-1).astype(int)
        draw.text((i[0], i[1] - offset), s, font=font, fill=(0, 0, 255))
    image_c = cv2.cvtColor(np.asarray(img_PIL), cv2.COLOR_RGB2BGR)
    cv2.imwrite(img_name01, image_c)
    
    
def draw_img(img, out, img_name01):
    text = out
    box_list = []
    for k in text.keys():
        if k in ('key', 'type'):
            continue
        elif len(text[k]['box']) > 0:
            box_list.append((k+':'+text[k]['text'], text[k]['box']))
    draw_pic(img, box_list, img_name01)

def box_take0(elem):
    return elem['points'][0][0]

def box_take1(elem):
    return elem['points'][0][1]

class CHART:
    def __init__(self, json_data):
        self.shapes = json_data['shapes']
        self.shapes.sort(key= box_take0)  # 按照左右排序
        self.shapes_u2d = self.shapes.copy()
        self.shapes_u2d.sort(key= box_take1)  # 按照上下排序
        self.imagePath = json_data['imagePath']
        self.imageHeight = json_data['imageHeight']
        self.imageWidth = json_data['imageWidth']
        self.title = {'text': '', 'box': [] }
        self.datasource = {'text': '', 'box': [] }

    def parse(self ):
        ## t1 = time.time()    
        self._parse_title()
        self._parse_datasource()
    
    ## title
    def _parse_title(self):
        is_title = False
        for shape in self.shapes:
            text = shape['label']
            box = shape['points']
            height = np.array(box)[:, 1].max()
            match01 = re.findall('^图[0-9]*', text)
            match02 = re.findall('^表[0-9]*', text)
            match03 = re.findall('^图表[0-9]*', text)
            if (len(match01) > 0 or len(match02) > 0 or len(match03) > 0) and height < self.imageHeight/4 and not is_title:
                self.title['text'] = text
                self.title['box'] = box
                is_title = True
                continue
            if len(self.title['box']) > 0 and is_title \
                and box[0][1] < self.title['box'][3][1] and box[3][1] > self.title['box'][0][1]:
                self.title['text'] = self.title['text'] + text
                x1 = min(self.title['box'][0][0], box[0][0])
                y1 = min(self.title['box'][0][1], box[0][1])
                x2 = max(self.title['box'][2][0], box[2][0])
                y2 = max(self.title['box'][2][1], box[2][1])
                self.title['box'] = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]

        if len(self.title['box']) == 0:
            text01 = ''
            box01 = []
            is_title = False
            h01 = 0
            for shape in self.shapes_u2d:
                text = shape['label']
                box = shape['points']
                height = np.array(box)[:,1].max()
                if height < self.imageHeight/4 and not is_title:
                    text01 = text
                    box01 = box.copy()
                    h01 = abs(box[2][1] - box[0][1])
                    is_title = True
                    continue
                if len(box01) > 0 and len(text01) > 0 and h01 > 0 and box[0][1] < box01[3][1] and box[3][1] > box01[0][1]:
                    box_mid = np.array(box)[:, 0].mean()
                    box01_mid = np.array(box01)[:, 0].mean()
                    if box_mid > box01_mid:
                        distance = box[0][0] - box01[1][0]
                        text02 = text01 + text
                    else:
                        distance = box01[0][0] - box[1][0]
                        text02 = text + text01
                    if distance < h01:
                        text01 = text02
                        x1 = min(box01[0][0], box[0][0])
                        y1 = min(box01[0][1], box[0][1])
                        x2 = max(box01[2][0], box[2][0])
                        y2 = max(box01[2][1], box[2][1])
                        box01 = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
                    else:
                        text01 = ''
                        box01 = []
                        break
            chin_length = len(''.join(re.findall('[\u4e00-\u9fa5]+', text01)))
            if chin_length > 0:
                self.title['text'] = text01
                self.title['box'] = box01

    ## datasource
    def _parse_datasource(self):
        is_datasource = False
        for shape in self.shapes:
            text = shape['label']
            box = shape['points']
            height = np.array(box)[:, 1].min()
            tmp01 = [i for i in ['资', '料', '来', '源'] if i in text]
            tmp02 = [i for i in ['数', '据', '来', '源'] if i in text]
            if ( len(tmp01) >= 3 or len(tmp02) >= 3 or '来源' in text ) and height > 0.75*self.imageHeight and not is_datasource:
                self.datasource['text'] = text
                self.datasource['box'] = box
                is_datasource = True
                continue                
            if len(self.datasource['box']) > 0 and is_datasource \
                and box[0][1] < self.datasource['box'][3][1] and box[3][1] > self.datasource['box'][0][1]:
                self.datasource['text'] = self.datasource['text'] + text
                x1 = min(self.datasource['box'][0][0], box[0][0])
                y1 = min(self.datasource['box'][0][1], box[0][1])
                x2 = max(self.datasource['box'][2][0], box[2][0])
                y2 = max(self.datasource['box'][2][1], box[2][1])
                self.datasource['box'] = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]            
    
    def pure_result(self):
        dict_result = { 'type': 'chart',
                        'key': { 'title': '表头',
                                 'datasource': '数据来源' },
                        'title': self.title,
                        'datasource': self.datasource  }
        return dict_result


if __name__ == '__main__':
    
    paths = ['./0/*.json',  './1/*.json',   './2/*.json',   './3/*.json' ]
    ###paths = ['./000/*.json' ]

    for path in paths:
        
        json_paths = glob(path)
        
        for json_path in tqdm(json_paths[:50000]):
            with open(json_path, 'r', encoding='utf-8' ) as f:
                json_data = json.load(f)
            
            chart = CHART(json_data)
            chart.parse()
            out = chart.pure_result()
            
            img_path = json_path.replace('.json', '.jpg')
            im = cv2.imread( img_path)
            
            img_name01 = './res' +img_path[2:]

            draw_img(im, out, img_name01)
    
    
    
    





